from math import sin, cos, pi
from stl import mesh
import numpy as np
import os
import tqdm
from scipy.interpolate import RectBivariateSpline

LABEL_BYTES = 984003
BYTES_PER_PIXEL = 2
LINES = 61440
LINE_SAMPLES = 184320
LATITUDE_RANGE = 2 / 3 * pi  # rad
LONGITUDE_RANGE = 2 * pi  # rad
scale_factor = 59.225
base_batch = 960
rate = 3
raw_reso = 59
tags = "sp"  # sphere
radius = 1737.4 * 1000

def sphere2plane(lat, lon):
    """
    Warning ! The unit of latitude and longitude is degree.
    """
    # 这里要加上一半的图像宽度，这是因为标注数据集中的角度是从-60度到+60度，起始位置不为零，因此需要修正。
    x = (pi / 3 - lat) * LINES / LATITUDE_RANGE
    # 这里不加一半的图像高度，这是因为标注数据集中的角度是从0到360度的，但是，图像的经度是从-180到+180度的，因此需要修正。
    y = lon * LINE_SAMPLES / LONGITUDE_RANGE
    return x, y


# 使用极坐标系解决问题是可行的
def create_sphere_segment(
    dem_image,
    row,
    col,
    radius,
    theta_start,
    delta_theta,
    phi_start,
    delta_phi,
    output_path,
    rate=2,
):
    # 生成球面片段的顶点
    vertices = []
    x = np.arange(0, dem_image.shape[0])
    y = np.arange(0, dem_image.shape[1])
    f = RectBivariateSpline(x, y, dem_image / scale_factor)
    num_steps = dem_image.shape[0] // rate

    for i in range(num_steps):
        for j in range(num_steps):
            # 纬度逐渐减小，经度逐渐增加
            theta = theta_start - i / (num_steps - 1) * delta_theta
            phi = phi_start + j / (num_steps - 1) * delta_phi
            xnew, ynew = sphere2plane(theta, phi)
            # 就是这里造成了接缝
            r = (
                f(
                    xnew - row * (dem_image.shape[0] - rate),
                    ynew - col * (dem_image.shape[1] - rate),
                )[0, 0]
                + radius / scale_factor
            ) / rate

            assert (
                xnew - row * (dem_image.shape[0] - rate)
                <= dem_image.shape[0] + rate + 1
                and xnew - row * (dem_image.shape[0] - rate) >= -rate - 1
            ) and (
                ynew - col * (dem_image.shape[1] - rate)
                <= dem_image.shape[1] + rate + 1
                and ynew - col * (dem_image.shape[1] - rate) >= -rate - 1
            ), f"Out of bound, {xnew -row * (dem_image.shape[0]-rate)}, {ynew - col * (dem_image.shape[1]-rate)}"

            x = r * cos(theta) * cos(phi)
            y = r * cos(theta) * sin(phi)
            z = r * sin(theta)
            vertices.append((x, y, z))

    # 生成三角形
    faces = []
    for i in range(num_steps - 1):
        for j in range(num_steps - 1):
            # 当前点的索引
            current = i * num_steps + j
            # 下一个theta值的点的索引
            next_theta = current + num_steps
            faces.append([current, next_theta, current + 1])
            faces.append([current + 1, next_theta, next_theta + 1])

    # 创建网格
    sphere_segment = mesh.Mesh(np.zeros(len(faces), dtype=mesh.Mesh.dtype))
    for i, f in enumerate(faces):
        for j in range(3):
            sphere_segment.vectors[i][j] = vertices[f[j]]

        # 保存为STL文件
    sphere_segment.save(output_path)


if __name__ == "__main__":
    path = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/Lunar_LRO_LOLAKaguya_DEMmerge_60N60S_512ppd.tif"
    output_dir = "/disk527/sdb1/a804_cbf/datasets/lunar_crater"

    num_batches = LINES // base_batch
    num_samples = LINE_SAMPLES // base_batch
    # 参数
    delta_theta = LATITUDE_RANGE / LINES * base_batch
    delta_phi = LONGITUDE_RANGE / LINE_SAMPLES * base_batch
    
    print(
        f"The basic batch of DEM is {base_batch}x{base_batch}. Each line contains {num_samples} batches. There are totally {num_batches*num_samples} basic batches."
    )
    with open(path, "rb") as f:
        # 跳过前984003的标签字节
        line = f.read(LABEL_BYTES)
        padding = [1, 1]
        for row in tqdm.tqdm(range(num_batches)):
            # 读取一个标准batch的数据，即960行，184320列，共计192个batch
            if row == 0:
                # 第一次读960+1行
                data = f.read(LINE_SAMPLES * BYTES_PER_PIXEL * (base_batch + rate))
                # 将数据转换为numpy数组
                data_array = np.frombuffer(data, dtype=np.int16).reshape(
                    base_batch + rate, LINE_SAMPLES
                )
                strip = data_array[-rate:, :]
                delta_theta_start = 0
            elif row == num_batches - 1:
                # 最后一次读960-1行，加上上一次保留的两行，相当于读取961行
                data = f.read(LINE_SAMPLES * BYTES_PER_PIXEL * (base_batch - rate))
                strip = data_array[-rate * 2 :, :]
                data_array = np.frombuffer(data, dtype=np.int16).reshape(
                    base_batch - rate, LINE_SAMPLES
                )
                data_array = np.vstack((strip, data_array))
                delta_theta_start = -delta_theta * rate * 2 / base_batch
            else:
                # 其余都读960行，加上上一次保留的一行，相当于都读取960+1行
                data = f.read(LINE_SAMPLES * BYTES_PER_PIXEL * base_batch)
                strip = data_array[-rate:, :]
                data_array = np.frombuffer(data, dtype=np.int16).reshape(
                    base_batch, LINE_SAMPLES
                )
                data_array = np.vstack((strip, data_array))
                delta_theta_start = 0
            if not os.path.exists(os.path.join(output_dir, "dem_int16")):
                os.makedirs(os.path.join(output_dir, "dem_int16"))
            if not os.path.exists(os.path.join(output_dir, "stl_models")):
                os.makedirs(os.path.join(output_dir, "stl_models"))
            if not os.path.exists(os.path.join(output_dir, "dem_int16", f"{row}")):
                os.makedirs(os.path.join(output_dir, "dem_int16", f"{row}"))
            if not os.path.exists(os.path.join(output_dir, "stl_models", f"{row}")):
                os.makedirs(os.path.join(output_dir, "stl_models", f"{row}"))
            for col in range(num_samples):
                # 用滑窗法填补每个块与块之间的缺失缝隙
                # 横向的滑窗，保证横向上都多一格，在最后一列时以前一格补后一格
                if col == num_samples - 1:
                    img_array = data_array[
                        :, col * base_batch - rate : (col + 1) * base_batch
                    ]
                    delta_phi_start = -delta_phi * rate / base_batch
                else:
                    # 重叠长度为1
                    img_array = data_array[
                        :, col * base_batch : (col + 1) * base_batch + rate
                    ]
                    delta_phi_start = 0
                # 创建球面片段
                theta_start = pi / 3 - (delta_theta * row + delta_theta_start)  # 北半球
                phi_start = delta_phi * col + delta_phi_start  # 东半球
                num_steps = base_batch // rate
                create_sphere_segment(
                    img_array,
                    row,
                    col,
                    radius,
                    theta_start,
                    delta_theta * (base_batch + rate) / base_batch,
                    phi_start,
                    delta_phi * (base_batch + rate) / base_batch,
                    os.path.join(
                        output_dir,
                        "stl_models",
                        f"{row}",
                        f"{tags}_{row}_{col}_{num_steps}_{raw_reso}m.stl",
                    ),
                    rate,
                )
                np.savez(
                    os.path.join(
                        output_dir,
                        "dem_int16",
                        f"{row}",
                        f"dem_{row}_{col}_{base_batch}.npz",
                    ),
                    image=img_array,
                )
